package com.bungarus.cluster;

import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.framework.recipes.cache.PathChildrenCache;
import org.apache.curator.retry.ExponentialBackoffRetry;
import org.apache.curator.utils.CloseableUtils;
import org.apache.curator.utils.ZKPaths;
import org.apache.log4j.Logger;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.data.Stat;
import org.springframework.util.StringUtils;

import java.util.HashMap;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

/**
 * Created by tang on 2018/11/14.
 */
public class ClusterManager {
    Logger logger = Logger.getLogger(ClusterManager.class);

    private static final String TAG_NEED_HIGH_PERFORMANCE = "busi";
    private static final String DEFAULT_HIGH_PERFORMANCE_NUMBER = "0";

    private static final String BASE_URL = "/com/bungarus/";
    private static final String BASE_CLUSTERS_TAG = "/clusters";
    private static final String BASE_CLUSTER_TAG = "/cluster";
    private static final String BASE_CUR_CLUSTER_TAG = "/cur_cluster";

    private static final String BASE_CLUSTER_PREFIX = "_";
    private static final String BASE_NODE_PREFIX = "node-";

    /**
     * zk cluster connection
     */
    private static CuratorFramework clusterConnection;
    /**
     * listen to the change of cluster's members
     */
    private Map<String, PathChildrenCache> pathCacheMap = new HashMap<>();
    /**
     * doing logic when change of cluster's members happen
     */
    private ClusterChangeListener clusterChangeListener = null;
    /**
     * whether zk ping scheduler started or not
     */
    private boolean started;
    /**
     * cluster tag
     */
    private String tag;

    public ClusterManager(String tag) {
        this(tag, ClusterConnection.DEFAULT_CONNECTION_STRING, ClusterConnection.DEFAULT_ZK_SESSION_TIMEOUT);
    }

    public ClusterManager(String tag, String connectStr) {
        this(tag, connectStr, ClusterConnection.DEFAULT_ZK_SESSION_TIMEOUT);
    }

    public ClusterManager(String tag, String connectStr, int sessionTimeout) {
        this.tag = tag;
        try {
            initCluster(connectStr, sessionTimeout);
        }
        catch (Exception e) {
            //throw new IllegalStateException(e);
            e.printStackTrace();
        }
    }

    private synchronized void initCluster(String connectStr, int sessionTimeout) throws Exception {
        if(null != clusterConnection){
            return;
        }
        /**
         * Notice that the ephemeral node will be lost at ZK when debugging with breakpoint-hanging in your IDE
         * and re-connecting timeout by the current strategy
         */
        clusterConnection = CuratorFrameworkFactory.builder()
                .connectString(connectStr)
                .connectionTimeoutMs(ClusterConnection.DEFAULT_ZK_CONNECTION_TIMEOUT)
                .sessionTimeoutMs(sessionTimeout)
                .retryPolicy(new ExponentialBackoffRetry(2000, 3)) // check session status every 1 second, if 3 times failed, make the session timeout
                .build();
        clusterConnection.start();

        initCluster0();
    }

    private void initCluster0() throws Exception {
        // if current cluster not exists, create it
        String currentCluster = getBaseCurrentClusterPath(this.tag);
        Stat exists = clusterConnection.checkExists().forPath(currentCluster);
        if(null == exists) {
            clusterConnection.create()
                    .creatingParentContainersIfNeeded()
                    .withMode(CreateMode.PERSISTENT)
                    .forPath(currentCluster);

            // create new cluster
            String clusterId = createCluster();

            // make the registered cluster as the current cluster
            if(clusterId.length() > 0){
                setCurrentCluster(clusterId);
            }
        }
    }

    /**
     * create a new cluster
     * @return the created cluster id prefixed with '_'
     * @throws Exception
     */
    public String createCluster() throws Exception {
        String clustersPath = getBaseClustersPath(this.tag);
        String clusterPath = getBaseClusterPath(this.tag);

        // create ".../clusters" path if not exists
        Stat exists = clusterConnection.checkExists().forPath(clustersPath);
        if(null == exists) {
            clusterConnection.create()
                    .creatingParentContainersIfNeeded()
                    .withMode(CreateMode.PERSISTENT)
                    .forPath(clustersPath);
        }
        // register the current cluster to path ".../clusters"
        String idPath = clusterConnection.create()
                //.withProtection()
                .withMode(CreateMode.PERSISTENT_SEQUENTIAL)
                .forPath(ZKPaths.makePath(clustersPath, BASE_CLUSTER_PREFIX)); // only use prefix here, and other part of the string will be generated automatically
        idPath = ZKPaths.getNodeFromPath(idPath); // get the auto-generated clusterId
        // create ".../cluster" path if not exists
        exists = clusterConnection.checkExists().forPath(clusterPath);
        if(null == exists) {
            clusterConnection.create()
                    .creatingParentContainersIfNeeded()
                    .withMode(CreateMode.PERSISTENT)
                    .forPath(clusterPath);
        }
        // register the current cluster to ".../cluster" path
        if(idPath.length() > 0){
            clusterConnection.create()
                    .creatingParentContainersIfNeeded()
                    .withMode(CreateMode.PERSISTENT)
                    .forPath(ZKPaths.makePath(clusterPath, idPath), "".getBytes("UTF-8"));
        }
        return idPath;
    }

    /**
     * make the cluster as current cluster
     * @param clusterId
     * @throws Exception
     */
    public void setCurrentCluster(String clusterId) throws Exception{
        clusterConnection.setData().forPath(getBaseCurrentClusterPath(this.tag), clusterId.getBytes("UTF-8"));
    }

    /**
     * set data for node
     * @param node
     * @throws Exception
     */
    public void setNodeData(NodeInfo node) throws Exception{
        String nodePath = ZKPaths.makePath(BASE_URL, node.getClusterBelongsTo(),
                BASE_CLUSTER_TAG,
                getFormattedClusterId(node.getClusterIDBelongsTo()),
                getFormattedNodeId(node.getNodeId()));

        String tagPortString = messagePorts2TagPortString(node.getMessagePorts());
        String nodeValue = buildNodeValue(nodeId2String(node.getNodeId()), node.getIp(),
                node.getProtocol(), node.getStatus(), tagPortString, node.isHighPerformance(), node.getExtraInfo());

        clusterConnection.setData().forPath(nodePath, nodeValue.getBytes("UTF-8"));
    }

    /**
     * register a node by auto-generate the nodeId to current cluster
     * @param node
     * @return NodeInfo
     * @throws RegisterNodeException
     */
    public NodeInfo registerNode(Node node) throws RegisterNodeException {
        //using zk client to register the specified node to be one member of a cluster
        long totalHighPerformanceNum = 0;
        long curHighPerformanceNum = 0;
        try {
            String currentClusterID = new String(clusterConnection.getData().forPath(getBaseCurrentClusterPath(this.tag)), "UTF-8");
            String nodePath = ZKPaths.makePath(getBaseClusterPath(node.getTag()), currentClusterID, BASE_NODE_PREFIX); // generate nodeId automaticaly
            String returnPath = clusterConnection.create()
                    //.withProtection()
                    .withMode(CreateMode.EPHEMERAL_SEQUENTIAL)
                    .forPath(nodePath);
            returnPath = ZKPaths.getNodeFromPath(returnPath);
            if(returnPath.length() <= BASE_NODE_PREFIX.length()){
                throw new RegisterNodeException();
            }

            String nodeID = returnPath.substring(BASE_NODE_PREFIX.length());  // get new generated nodeId

            if(node.getTag().equalsIgnoreCase(TAG_NEED_HIGH_PERFORMANCE)) {
                //get number of high performance in current cluster
                try {
                    totalHighPerformanceNum = Integer.parseInt(
                            this.getHighPerformanceString(node, currentClusterID));
                    List<NodeInfo> registeredNodes = this.nodes(node.getTag(), Integer.parseInt(currentClusterID.substring(1)));
                    if(null != registeredNodes && registeredNodes.size() > 0) {
                        curHighPerformanceNum = registeredNodes.stream()
                                .filter(nodeInfo -> nodeInfo != null)
                                .filter(NodeInfo::isHighPerformance).count();
                    }
                }
                catch (Exception e) {
                    //error happens and ignores the registration of node processing large groups
                    e.printStackTrace();
                }
            }
            boolean highPerformance = false;
            if(totalHighPerformanceNum > 0 && curHighPerformanceNum < totalHighPerformanceNum) {
                highPerformance = true;
            }

            //pattern-> {nodeId}:{ip}:{protocol}:{status}:{tagPortString}:{highPerformance}
            String nodeValue = buildNodeValue(nodeID, node.getHostAddress(), node.availableProtocol(),
                    Node.STATUS_INITIALIZING, node.tagPortString4Connection(), highPerformance, node.getExtraInfo());

            // set data for the node
            clusterConnection.setData().forPath(nodePath + nodeID, nodeValue.getBytes("UTF-8"));

            NodeInfo nodeInfo = new NodeInfo();
            nodeInfo.setNodeId(Integer.parseInt(nodeID));
            nodeInfo.setClusterIDBelongsTo(Integer.parseInt(currentClusterID.substring(1)));
            nodeInfo.setClusterBelongsTo(node.getTag());
            nodeInfo.setProtocol(node.availableProtocol());
            nodeInfo.setStatus(Node.STATUS_INITIALIZING);
            nodeInfo.setIp(node.getHostAddress());
            nodeInfo.setHighPerformance(highPerformance);
            nodeInfo.setExtraInfo(node.getExtraInfo());

            Map<String, Integer> map = tagPortString2MessagePorts(node.tagPortString4Connection());
            if (map.size() > 0) {
                nodeInfo.setMessagePorts(map);
            }
            logger.info("=================== registered nodeInfo, cluster: " + node.getTag() + " detail: " + nodeInfo.toString());
            return nodeInfo;
        }
        catch (Exception e) {
            throw new RegisterNodeException(e);
        }
    }

    private boolean tagHighPerformanceToNode(int times, Node node, String currentClusterID) throws Exception {
        if(times >= 3) {
            return false;
        }
        try {
            Stat stat = new Stat();
            String clusterHighPerformanceDataPath = ZKPaths.makePath(getBaseClusterPath(node.getTag()), currentClusterID);
            byte[] hp = clusterConnection.getData().storingStatIn(stat).forPath(clusterHighPerformanceDataPath);
            String[] v = StringUtils.tokenizeToStringArray(new String(hp, "UTF-8"), NodeInfo.SEPARATOR);
            final int total = Integer.parseInt(v[0]);
            int cur = Integer.parseInt(v[1]);
            if(cur < total) {
                cur++;
                clusterConnection.setData()
                        .withVersion(stat.getVersion())
                        .forPath(clusterHighPerformanceDataPath, (total + NodeInfo.SEPARATOR + cur).getBytes("UTF-8"));
            }
        }
        catch (Exception e) {
            e.printStackTrace();
            //failed retry
            Thread.sleep(1000);
            times++;
            tagHighPerformanceToNode(times, node, currentClusterID);
        }
        return false;
    }

    private String getHighPerformanceString(Node node, String currentClusterID) throws Exception {
        String clusterHighPerformanceDataPath = ZKPaths.makePath(getBaseClusterPath(node.getTag()), currentClusterID);
        Stat stat = new Stat();
        byte[] hp = clusterConnection.getData().storingStatIn(stat).forPath(clusterHighPerformanceDataPath);
        String ret = new String(hp, "UTF-8");
        if(ret.length() <= 0) {
            try {
                clusterConnection.setData()
                        .withVersion(stat.getVersion())
                        .forPath(clusterHighPerformanceDataPath, DEFAULT_HIGH_PERFORMANCE_NUMBER.getBytes("UTF-8"));
            }
            catch (Exception e) {
                e.printStackTrace();
            }
            finally {
                ret = DEFAULT_HIGH_PERFORMANCE_NUMBER;
            }
        }
        return ret;
    }

    public List<ClusterInfo> clusters() throws Exception {
        return clusters(this.tag);
    }

    /**
     * get clusters by tag
     * @param clusterTag
     * @return List<ClusterInfo>
     * @throws Exception
     */
    public List<ClusterInfo> clusters(String clusterTag) throws Exception {
        String clustersPath = getBaseClustersPath(clusterTag);
        List<String>clusterChildren = clusterConnection.getChildren().forPath(clustersPath);
        return clusterChildren.stream()
                .map(clusterPath -> {
                    ClusterInfo clusterInfo = new ClusterInfo();
                    clusterInfo.setTag(clusterTag);
                    clusterInfo.setClusterID(Integer.valueOf(clusterPath.substring(1)));
                    return clusterInfo;
                })
                .sorted((info1, info2) -> {
                    if (info1.getClusterID().equals(info2.getClusterID())) {
                        return 0;
                    }
                    else {
                        return info1.getClusterID() > info2.getClusterID() ? 1 : -1;
                    }
                })
                .collect(Collectors.toList());
    }


    public List<NodeInfo> nodes(int clusterId) throws Exception {
        return nodes(this.tag, clusterId);
    }

    /**
     * get node list by cluster from a tag
     * @param tag
     * @param clusterId
     * @return List<NodeInfo>
     * @throws Exception
     */
    public List<NodeInfo> nodes(String tag, int clusterId) throws Exception {
        String clusterPath = ZKPaths.makePath(getBaseClusterPath(tag), getFormattedClusterId(clusterId));
        List<String> nodePaths = clusterConnection.getChildren().forPath(clusterPath);
        return nodePaths.stream()
                .map(nodePath -> {
                    try {
                        //node data format: nodeId:ip:protocol:status:(tag1|port1,tag2|port2):high_performance
                        String nodeDataStr = new String(clusterConnection.getData().forPath(ZKPaths.makePath(clusterPath, nodePath)), "UTF-8");
                        return buildNodeInfo(tag, clusterId, nodeDataStr);
                    }
                    catch (Exception e) {
                        e.printStackTrace();
                    }
                    return null;
                })
                .collect(Collectors.toList());
    }

    /**
     * build node info from node value string
     * @param clusterId
     * @param nodeDataStr
     * @return
     */
    private NodeInfo buildNodeInfo(String tag, int clusterId, String nodeDataStr){
        String[] nodeData = StringUtils.tokenizeToStringArray(nodeDataStr, NodeInfo.SEPARATOR);
        if (nodeData.length == 7) {

            NodeInfo nodeInfo = new NodeInfo();
            nodeInfo.setNodeId(Integer.parseInt(nodeData[0]));
            nodeInfo.setClusterIDBelongsTo(clusterId);
            nodeInfo.setClusterBelongsTo(tag);
            nodeInfo.setIp(nodeData[1]);
            nodeInfo.setProtocol(Integer.parseInt(nodeData[2]));
            nodeInfo.setStatus(Integer.parseInt(nodeData[3]));

            String tagPortString = nodeData[4];
            Map<String, Integer> map = tagPortString2MessagePorts(tagPortString);
            if (map.size() > 0) {
                nodeInfo.setMessagePorts(map);
            }
            nodeInfo.setHighPerformance(
                    Integer.parseInt(nodeData[5]) > NodeInfo.NORMAL_PERFORMANCE
            );
            nodeInfo.setExtraInfo(nodeData[6]);
            return nodeInfo;
        }
        return null;
    }

    /**
     * build node value
     * pattern -> {nodeId}:{ip}:{protocol}:{status}:({tag}|{port},...)
     * @param nodeID
     * @param ip
     * @param status  node status
     * @param tagPortString
     * @param highPerformance
     * @param extraInfo  extra info
     */
    private String buildNodeValue(String nodeID, String ip, int protocol, int status, String tagPortString, boolean highPerformance, String extraInfo){
        StringBuilder sb = new StringBuilder(nodeID);
        sb.append(NodeInfo.SEPARATOR);
        sb.append(ip);
        sb.append(NodeInfo.SEPARATOR);
        sb.append(protocol);
        sb.append(NodeInfo.SEPARATOR);
        sb.append(status);
        sb.append(NodeInfo.SEPARATOR);
        sb.append(tagPortString);
        sb.append(NodeInfo.SEPARATOR);
        if(highPerformance) {
            sb.append(NodeInfo.HIGH_PERFORMANCE);
        }
        else {
            sb.append(NodeInfo.NORMAL_PERFORMANCE);
        }
        sb.append(NodeInfo.SEPARATOR);
        sb.append(extraInfo);
        return sb.toString();
    }

    /**
     * convert tagPortString into Map
     * @param tagPortString
     * @return Map<String, Integer>
     */
    private Map<String, Integer> tagPortString2MessagePorts(String tagPortString) {
        Map<String, Integer> map = new HashMap<>(0);;
        if (tagPortString.length() >= 2) { // 2 stands for including only ()
            tagPortString = tagPortString.substring(1, tagPortString.length() - 1);
            String[] tagPortStrs = StringUtils.tokenizeToStringArray(tagPortString, ",");
            if (null != tagPortStrs && tagPortStrs.length > 0) {
                for (String tagPort : tagPortStrs) { //tag1|port1
                    String[] kv = StringUtils.split(tagPort, NodeInfo.TAG_PORT_SEPARATOR);
                    if (null != kv && kv.length >= 2) {
                        map.put(kv[0], Integer.valueOf(kv[1]));
                    }
                }
            }
        }
        return map;
    }

    /**
     * convert Map of message ports into tagPortString
     * @param messagePorts
     * @return String
     */
    private String messagePorts2TagPortString(Map<String, Integer> messagePorts) {
        if(null != messagePorts){
            Iterator<String> tags = messagePorts.keySet().iterator();
            // string format: "({tag}|{port})"
            StringBuilder sb = new StringBuilder("(");
            while(tags.hasNext()){
                String tag = tags.next();
                sb.append(tag);
                sb.append(NodeInfo.TAG_PORT_SEPARATOR);
                sb.append(messagePorts.get(tag));
                sb.append(",");
            }
            if(sb.length() > 1) { // remove comma
                String result = sb.substring(0, sb.length() - 1);
                return result + ")";
            }
        }
        return "";
    }

    public void close() {
        this.pathCacheMap.keySet().forEach(key -> {
            CloseableUtils.closeQuietly(pathCacheMap.get(key));
        });
        CloseableUtils.closeQuietly(clusterConnection);
    }

    public void addClusterChangeListener(ClusterChangeListener clusterChangeListener) {
        this.clusterChangeListener = clusterChangeListener;
        this.clusterChangeListener.getClusterTags().stream().forEach(clusterTag -> {
            String clustersPath = getBaseClustersPath(clusterTag);
            List<String> children = null;
            try {
                children = clusterConnection.getChildren().forPath(clustersPath);
            }
            catch (Exception e) {
                throw new RuntimeException(e);
            }
            children.stream().forEach(cluster -> {
                String clusterPath = ZKPaths.makePath(getBaseClusterPath(clusterTag), cluster);
                if(null == this.pathCacheMap.get(clusterPath)) {
                    PathChildrenCache pathCache = new PathChildrenCache(clusterConnection, clusterPath, true);
                    pathCache.getListenable().addListener((curatorFramework, pathChildrenCacheEvent) -> {
                        switch(pathChildrenCacheEvent.getType()) {
                            case CHILD_ADDED: {
                                String nodeDataStr = new String(pathChildrenCacheEvent.getData().getData(), "UTF-8");
                                NodeInfo node = buildNodeInfo(clusterTag, clusterString2Id(cluster), nodeDataStr);
                                if (null != node) {
                                    logger.info("Node added -> " + clusterTag + ":" + nodeId2String(node.getNodeId()));
                                    this.clusterChangeListener.onClusterNodeAdded(node);
                                }
                                break;
                            }
                            case CHILD_UPDATED: {
                                String nodeDataStr = new String(pathChildrenCacheEvent.getData().getData(), "UTF-8");
                                if(null != nodeDataStr && nodeDataStr.length() > 0) {
                                    NodeInfo node = buildNodeInfo(clusterTag, clusterString2Id(cluster), nodeDataStr);
                                    if(null != node) {
                                        logger.info("Node updated -> tag:" + clusterTag + "|id:" + nodeId2String(node.getNodeId()) + "|status:" + node.getStatus());
                                        if (node.getStatus() == Node.STATUS_RUNNING) {
                                            this.clusterChangeListener.onClusterNodeUpdated(node);
                                        }
                                    }
                                }
                                break;
                            }
                            case CHILD_REMOVED: {
                                String nodeDataStr = new String(pathChildrenCacheEvent.getData().getData(), "UTF-8");
                                NodeInfo node = buildNodeInfo(clusterTag, clusterString2Id(cluster), nodeDataStr);
                                if (null != node) {
                                    logger.info("Node removed -> " + clusterTag + ":" + node.getNodeId());
                                    this.clusterChangeListener.onClusterNodeRemoved(node);
                                }
                                break;
                            }
                            default:
                                break;
                        }
                    });
                    try {
                        pathCache.start();
                        this.pathCacheMap.put(clusterPath, pathCache);
                    }
                    catch (Exception e) {
                        throw new RuntimeException(e);
                    }
                }
            });
        });
    }


    public String getBaseClusterPath(String  tag){
        return ZKPaths.makePath(BASE_URL, tag, BASE_CLUSTER_TAG);
    }

    public String getBaseClustersPath(String tag){
        return ZKPaths.makePath(BASE_URL, tag, BASE_CLUSTERS_TAG);
    }

    public String getBaseCurrentClusterPath(String tag){
        return ZKPaths.makePath(BASE_URL, tag, BASE_CUR_CLUSTER_TAG);
    }

    public String clusterId2String(int clusterId){
        return String.format("%010d", clusterId);
    }

    public int clusterString2Id(String clusterId){
        if(clusterId.indexOf(BASE_CLUSTER_PREFIX) == 0){
            clusterId = clusterId.substring(BASE_CLUSTER_PREFIX.length());
        }
        return Integer.parseInt(clusterId);
    }

    public String getFormattedClusterId(int clusterId){
        return BASE_CLUSTER_PREFIX + clusterId2String(clusterId);
    }

    public String nodeId2String(int nodeId){
        return String.format("%010d", nodeId);
    }

    public int nodeString2Id(String nodeId){
        if(nodeId.indexOf(BASE_NODE_PREFIX) == 0){
            nodeId = nodeId.substring(BASE_NODE_PREFIX.length());
        }
        return Integer.parseInt(nodeId);
    }

    public String getFormattedNodeId(int nodeId){
        return BASE_NODE_PREFIX + nodeId2String(nodeId);
    }

}
